# pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 -i https://pypi.tuna.tsinghua.edu.cn/simple

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
# from torch.optim.lr_scheduler import ReduceLROnPlateau

# --------------------- 检查GPU设备 ---------------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# print(f"Using device: {device}")

# --------------------- 数据读取和预处理 ---------------------
data = pd.read_csv("E:/GraduateDesign/LinearUse_processed.csv")

# #matplot字体选择
# plt.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体
# plt.rcParams['axes.unicode_minus'] = False    # 解决负号显示问题
# # 数据分析
# # print("\n===== 目标变量(rrr)分布分析 =====")
# # print(data['rrr'].describe())
# plt.figure(figsize=(10,4))
# plt.subplot(1,2,1)
# plt.hist(data['rrr'], bins=50)
# plt.title("原始数据分布")
# plt.subplot(1,2,2)
# plt.boxplot(data['rrr'], vert=False)
# plt.title("箱线图")
# plt.show()


X = data.drop(columns=['rrr']).values
y = data['rrr'].values.reshape(-1, 1)

# 标准化数据
scaler_X = StandardScaler()

X_scaled = scaler_X.fit_transform(X)

scaler_y = MinMaxScaler(feature_range=(0.1, 1))

y_scaled = scaler_y.fit_transform(y)

# 转换为PyTorch张量并发送到设备
X_reshaped = X_scaled.reshape(-1, 1, X_scaled.shape[1])
X_tensor = torch.tensor(X_reshaped, dtype=torch.float32).to(device)  # 直接发送到设备
y_tensor = torch.tensor(y_scaled, dtype=torch.float32).to(device)

# 划分训练集和测试集（现在直接在GPU上划分）
X_train, X_test, y_train, y_test = train_test_split(
    X_tensor, y_tensor, test_size=0.2, random_state=42, shuffle=True
)

# scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5)
# --------------------- 定义Transformer模型 ---------------------
class TransformerRegressor(nn.Module):
    def __init__(self, input_dim, output_dim, num_heads=4, dim_feedforward=64, num_layers=2):
        super().__init__()
        self.encoder_layer = nn.TransformerEncoderLayer(
            d_model=input_dim,
            nhead=num_heads,
            dim_feedforward=dim_feedforward,
            batch_first=True
        )
        self.transformer_encoder = nn.TransformerEncoder(
            self.encoder_layer,
            num_layers=num_layers
        )
        self.fc = nn.Linear(input_dim, output_dim)

        self.fc = nn.Sequential(
            nn.Linear(input_dim, output_dim),
            nn.ReLU()  # 确保输出非负
        )
    def forward(self, x):
        out = self.transformer_encoder(x)
        return self.fc(out[:, -1, :])


# 初始化模型并发送到GPU
model = TransformerRegressor(input_dim=X_train.shape[2], output_dim=1).to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()

# --------------------- 训练模型 ---------------------
train_losses = []
val_losses = []
num_epochs = 900
batch_size = 320

# 创建DataLoader时保持数据在GPU上
train_dataset = torch.utils.data.TensorDataset(X_train, y_train)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

for epoch in range(num_epochs):
    model.train()
    epoch_loss = 0
    for X_batch, y_batch in train_loader:
        # 数据已经在GPU上，无需额外转移
        optimizer.zero_grad()
        y_pred = model(X_batch)
        loss = criterion(y_pred, y_batch)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()

    # 验证损失（使用部分训练数据）
    model.eval()
    with torch.no_grad():
        val_pred = model(X_train[:len(X_test)])  # 使用相同大小的数据验证
        val_loss = criterion(val_pred, y_train[:len(y_test)])

    train_losses.append(epoch_loss / len(train_loader))
    val_losses.append(val_loss.item())

    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch + 1}/{num_epochs}, Train Loss: {train_losses[-1]:.4f}, Val Loss: {val_losses[-1]:.4f}")

# # --------------------- 评估测试集 ---------------------
# model.eval()
# with torch.no_grad():
#     y_pred_scaled = model(X_test).cpu().numpy()  # 移回CPU处理
#
# # 反标准化
# y_test_original = scaler_y.inverse_transform(y_test.cpu().numpy())  # 移回CPU处理
# y_pred = scaler_y.inverse_transform(y_pred_scaled)
#
# # 展平数据（确保形状匹配）
# y_test_flat = y_test_original.flatten()
# y_pred_flat = y_pred.flatten()
#
# # --------------------- 计算所有评估指标 ---------------------
# # 处理潜在的零值（如果数据中存在零）
# epsilon = 1e-10  # 防止除零的小量
# y_test_nonzero = np.where(y_test_flat == 0, epsilon, y_test_flat)
#
# # 计算各项指标
# rmse = mean_squared_error(y_test_flat, y_pred_flat, squared=False)
# mae = mean_absolute_error(y_test_flat, y_pred_flat)
# mape = np.mean(np.abs((y_test_flat - y_pred_flat) / y_test_nonzero)) * 100
# r2 = r2_score(y_test_flat, y_pred_flat)
#
# # 输出结果
# print("\n================ 评估结果 ================")
# print(f"RMSE: {rmse:.4f}")
# print(f"MAE: {mae:.4f}")
# print(f"MAPE: {mape:.4f}%")
# print(f"R²: {r2:.4f}")
#
# # 如果存在零值给出警告
# if np.any(y_test_flat == 0):
#     print("\n警告：原始数据中存在零值，MAPE计算可能失真！")

# --------------------- 评估测试集 ---------------------
model.eval()
with torch.no_grad():
    y_pred_scaled = model(X_test).cpu().numpy()

# 反标准化
y_test_original = scaler_y.inverse_transform(y_test.cpu().numpy())
y_pred = scaler_y.inverse_transform(y_pred_scaled)

# 展平数据
y_test_flat = y_test_original.flatten()
y_pred_flat = y_pred.flatten()

# # 诊断输出
# print("\n===== 数据验证 =====")
# print(f"实际值范围: [{y_test_flat.min():.4f}, {y_test_flat.max():.4f}]")
# print(f"预测值范围: [{y_pred_flat.min():.4f}, {y_pred_flat.max():.4f}]")
# print("实际值零/负数量:", np.sum(y_test_flat <= 0))

# 仅当存在有效实际值时计算MAPE
valid_indices = y_test_flat > 0
if np.sum(valid_indices) > 0:
    y_test_valid = y_test_flat[valid_indices]
    y_pred_valid = y_pred_flat[valid_indices]
    mape = np.mean(np.abs((y_test_valid - y_pred_valid) / y_test_valid)) * 100
else:
    mape = float('nan')

# 计算SMAPE
smape = np.mean(2 * np.abs(y_pred_flat - y_test_flat) / (np.abs(y_test_flat) + np.abs(y_pred_flat) + 1e-10)) * 100

# 其他指标
rmse = mean_squared_error(y_test_flat, y_pred_flat, squared=False)
mae = mean_absolute_error(y_test_flat, y_pred_flat)
r2 = r2_score(y_test_flat, y_pred_flat)

print("\n================ 评估结果 ================")
print(f"RMSE: {rmse:.4f}")
print(f"MAE: {mae:.4f}")
# if not np.isnan(mape):
#    print(f"MAPE: {mape:.4f}%")
print(f"MAPE: {smape:.4f}%")
print(f"R²: {r2:.4f}")



#--------------------- 可视化训练曲线（可选） ---------------------

plt.plot(train_losses, label='Training Loss')
plt.plot(val_losses, label='Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('MSE Loss')
plt.legend()
plt.show()